// expression divergence analysis
// by Latent Dirichlet Allocation

#include <boost/ptr_container/ptr_vector.hpp>
#include <vector>
#include <fstream>
#include <iostream>
#include <string>
#include <sstream>

#include "ted.hh"
#include "expression.hh"
#include "latent.hh"
#include "lda.hh"

#include <numeric>
#include <algorithm>

typedef geneset_t<expr_pair_t> ep_geneset_t;
typedef boost::ptr_vector<ep_geneset_t> ep_geneset_vec_t;
typedef boost::ptr_vector<expr_pair_set_t> pair_set_vec_t;

void output_topics(std::string& output,
    boost::ptr_vector<func_expr_div_t>& topics)
{
  std::ofstream ofs_mean((output + ".mean").c_str(), std::ios::out);
  std::ofstream ofs_prec((output + ".prec").c_str(), std::ios::out);
  std::ofstream ofs_sd((output + ".sd").c_str(), std::ios::out);

  for (size_t k = 0; k < topics.size(); ++k)
  {
    const func_expr_div_t::vec_t& mean = topics[k].mean_vec();
    const func_expr_div_t::vec_t& var = topics[k].var_vec();
    const func_expr_div_t::vec_t& prec = topics[k].prec_vec();

    ofs_mean << k << "\t";
    vec_out(ofs_mean, mean);
    ofs_mean << std::endl;
    ofs_sd << k << "\t";
    vec_out(ofs_sd, var, sqrt_op_t());
    ofs_sd << std::endl;
    ofs_prec << k << "\t";
    vec_out(ofs_prec, prec);
    ofs_prec << std::endl;
  }

  ofs_mean.close();
  ofs_prec.close();
  ofs_sd.close();
}

struct prog_args_t
{

  explicit prog_args_t() :
      g_file(""), x_file(""), y_file(""), K(2), // required args
      output("output"), bmin(-2.), bmax(2.),    // beta in [bmin, bmax]
      repeat(1), tol(1e-3), max_iter(100), sample(100), // options
      minibatch(1.)
  {
  }

  std::string g_file;
  std::string x_file;
  std::string y_file;
  size_t K;
  std::string output;
  double bmin;
  double bmax;
  size_t repeat;
  double tol;
  size_t max_iter;
  size_t sample;
  size_t minibatch;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -g G.txt -x X.txt -y Y.txt -k K -o out" << std::endl;
  std::cerr << "         -bmin bmin -bmax bmax" << std::endl;
  std::cerr << "         -i iter -r repeat -t tol -s sample" << std::endl;
#ifdef ONLINE
  std::cerr << "         -minibatch size" << std::endl;
#endif
  std::cerr << std::endl;
  std::cerr
      << " G.txt  geneset / group file (each line contains row-indexes of X,Y)"
      << std::endl;
  std::cerr << " X.txt  control expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " Y.txt  case expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " K      number of topics (default: " << args.K << ")"
      << std::endl;
  std::cerr << " out    output files' header (default: " << args.output << ")"
      << std::endl;
  std::cerr << " bmin   min value of beta (default: " << args.bmin << ")"
      << std::endl;
  std::cerr << " bmax   max value of beta (default: " << args.bmax << ")"
      << std::endl;
  std::cerr << " iter   max number of iterations (default: " << args.max_iter
      << ")" << std::endl;
  std::cerr << " tol    convergence condition (default: " << args.tol << ")"
      << std::endl;
  std::cerr << " sample number of samples (default: " << args.sample << ")"
      << std::endl;
  std::cerr << " repeat number of repeated fits (default: " << args.repeat
      << ")" << std::endl;
#ifdef ONLINE
  std::cerr << " minibatch size of mini-batch (default: " << args.minibatch
      << ")" << std::endl;
#endif
  std::cerr << std::endl;
  std::cerr << ".argmax .genesets .mean .prec .sd" << std::endl;
  std::cerr << std::endl;
}

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-g" && ++j < argc)
    {
      args_out.g_file = argv[j];
    }
    else if (curr == "-x" && ++j < argc)
    {
      args_out.x_file = argv[j];
    }
    else if (curr == "-y" && ++j < argc)
    {
      args_out.y_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-i" && ++j < argc)
    {
      args_out.max_iter = get_param<size_t>(argv[j], "parsing iter");
    }
    else if (curr == "-s" && ++j < argc)
    {
      args_out.sample = get_param<size_t>(argv[j], "parsing #samples");
    }
    else if (curr == "-k" && ++j < argc)
    {
      args_out.K = get_param<size_t>(argv[j], "parsing K");
    }
    else if (curr == "-r" && ++j < argc)
    {
      args_out.repeat = get_param<size_t>(argv[j], "parsing repeat");
    }
    else if (curr == "-minibatch" && ++j < argc)
    {
      args_out.minibatch = get_param<size_t>(argv[j], "parsing minibatch");
    }
    else if (curr == "-t" && ++j < argc)
    {
      args_out.tol = get_param<double>(argv[j], "parsing tol");
    }
    else if (curr == "-bmin" && ++j < argc)
    {
      args_out.bmin = get_param<double>(argv[j], "parsing beta min");
    }
    else if (curr == "-bmax" && ++j < argc)
    {
      args_out.bmax = get_param<double>(argv[j], "parsing beta max");
    }
  }

  if (args_out.tol <= 0)
    return false;
  if (args_out.g_file.size() == 0 || args_out.x_file.size() == 0
      || args_out.y_file.size() == 0)
    return false;
  if (args_out.bmin >= args_out.bmax)
    return false;
  if (args_out.K < 2)
    return false;
  return true;
}

int main(const int argc, const char* argv[])
{

  prog_args_t args;

  if (!parse_args(argc, argv, args))
  {
    print_help(argv[0], args);
    return 1;
  }

  // read expression data
  expr_pair_vec_t pair_data;
  size_t n, p;
  std::string& x = args.x_file;
  std::string& y = args.y_file;
  assert_msg(read_expr_pair_vec(x.c_str(), y.c_str(), pair_data, n, p),
      "cannot read expression pairs");

  // read geneset data
  std::string& g = args.g_file;
  ep_geneset_vec_t genesets;
  double a0 = 1. / ((double) args.K);
  const size_t K = args.K;
  assert_msg(read_genesets<expr_pair_t>(g.c_str(), pair_data, a0, K, genesets),
      "cannot read genesets");

#ifdef ONLINE
  // required for stochastic online update
  std::vector<int> geneset_order;
  for (size_t j = 0; j < genesets.size(); ++j)
    geneset_order.push_back(j);
#endif

  double llik_max = 0.;
  for (size_t r = 1; r <= args.repeat; ++r)
  {
    // initialize K topics
    boost::ptr_vector<func_expr_div_t> topics;
    for (size_t k = 0; k < K; ++k)
      topics.push_back(new func_expr_div_t(p, args.bmin, args.bmax));

    boost::ptr_vector<expr_pair_set_t> pairsets;
    for (size_t k = 0; k < K; ++k)
      pairsets.push_back(new expr_pair_set_t);

    // initialization by random seeding of genes to K topics
    random_seeding<expr_pair_t, func_expr_div_t>(pair_data, pairsets, topics,
        K);

    // update function by Gibbs sampling
    gibbs_latent_update_t<expr_pair_t, func_expr_div_t> gibbs_update(topics,
        args.sample);

    const double t0 = 1.;
    const double rate_discount = 0.75;

#ifdef ONLINE

    for (size_t eb_iter = 1; eb_iter <= args.max_iter; ++eb_iter)
    {

      std::random_shuffle(geneset_order.begin(), geneset_order.end());

      const size_t n = genesets.size();
      const size_t m = std::min(n, args.minibatch);
      const double D = genesets.size();

      for (size_t s = 0; s < n; s += m)
      {
        double t = 1. + ((double) s);
        double rate = std::pow(t + t0, -rate_discount);

        for (size_t k = 0; k < pairsets.size(); ++k)
          pairsets[k].clear();

        double S = 0.;
        for (size_t j = s; j < std::min(s + m, n); ++j)
        {
          S++;
          geneset_t<expr_pair_t>& geneset = genesets[geneset_order[j % n]];
          gibbs_update(geneset);
          // construct training set
          for (size_t i = 0; i < geneset.n; ++i)
          {
            const latent_matrix_t::vec_t& z_i = geneset.Z(i);
            for (size_t k = 0; k < geneset.K; ++k)
              if (z_i.at(k) > 0)
                pairsets[k].insert(geneset.data.at(i), z_i.at(k));
          }
        }

        // actually update topics
        double delta = 0.;
        for (size_t k = 0; k < K; ++k)
          delta += topics[k].update(pairsets[k], rate, D / S);
        delta /= ((double) K);

        TLOG(std::setw(15) << rate << std::setw(15) << delta);

        if (delta < args.tol)
          break;
      }

      // empirical bayes
      double eb_rate = std::pow((double) eb_iter + t0, -rate_discount);
      double eb_delta = empirical_bayes(topics, eb_rate);

      TLOG("static:" << std::setw(10) << topics[0].get_penalty_static() << //
          ", kinetic:" << std::setw(10) << topics[0].get_penalty_kinetic());

      if (eb_delta < args.tol)
        break;
    }
#else
    for (size_t eb_iter = 1; eb_iter <= args.max_iter; ++eb_iter)
    {

      for (double t = 1; t <= args.max_iter; ++t)
      {
        double rate = std::pow(t + t0, -rate_discount);

        std::for_each(genesets.begin(), genesets.end(), gibbs_update);

        construct_training_sets(genesets, pairsets);

        double delta = 0.;
        for (size_t k = 0; k < K; ++k)
        delta += topics[k].update(pairsets[k], rate);
        delta /= ((double) K);

        TLOG(std::setw(15) << rate << std::setw(15) << delta);

        if (delta < args.tol)
        break;
      }

      double eb_rate = std::pow( (double) eb_iter + t0, - rate_discount );
      double eb_delta = empirical_bayes(topics, eb_rate);

      TLOG("static:" << std::setw(10) << topics[0].get_penalty_static() << //
          ", kinetic:" << std::setw(10) << topics[0].get_penalty_kinetic());

      if (eb_delta < args.tol)
      break;
    }
#endif

    // final latent update
    std::for_each(genesets.begin(), genesets.end(), gibbs_update);
    double llik = likelihood(genesets);

    TLOG("log-likelihood = " << llik);

    if (r == 1 || llik > llik_max)
    {
      llik_max = llik;
      output_topics(args.output, topics);
      output_argmax(args.output, pair_data, topics);
      output_genesets(args.output, genesets);
    }
  }

  return 0;
}

